#include "robot_one_navigation/A_Star/Astar.hpp"

namespace astar_planner
{
    Astar::Astar()
    {
        allow_unknown_ = false;
        motion_model_ = MotionModel::MOORE;
    }
    Astar::~Astar()
    {
    }

    // 清除openlist
    void Astar::clear_queue()
    {
        NodeQueue q;
        std::swap(open_list_, q);
    }

    // 设置代价地图
    void Astar::set_costmap(nav2_costmap_2d::Costmap2D *costmap)
    {
        costmap_ = costmap;
        xs_ = costmap_->getSizeInCellsX();
        ys_ = costmap_->getSizeInCellsY();
        ns_ = xs_ * ys_;

        clear_queue();

        graph_.clear();
        graph_.reserve(ns_);

        // 初始化所有栅格作为node
        for (int j = 0; j < ys_; ++j)
        {
            for (int i = 0; i < xs_; ++i)
            {
                float v = OCCUPIED;
                v = interpret_costmap(i, j, costmap_);
                graph_.emplace_back(get_index(i, j), i, j, v);
            }
        }
    }

    float Astar::interpret_costmap(int i, int j, nav2_costmap_2d::Costmap2D *costmap)
    {
        float cost = static_cast<float>(costmap->getCost(i, j));
        if (allow_unknown_ && cost == UNKNOWN)
        {
            return MAX_NON_OBSTACLE;
        }
        else if (!allow_unknown_ && cost == UNKNOWN)
        {
            return OCCUPIED;
        }
        else if (cost == INSCRIBED)
        {
            return OCCUPIED;
        }
        else
        {
            return cost;
        }
    }

    /*
    计算启发值h_x
    h = min(dx,dy)*√2 + |dx–dy|*1
      = min(dx,dy)*√2 + (max(dx,dy)–min(dx,dy))*1
      = min(dx,dy)*(√2 – 1) + max(dx,dy)
      = min(dx,dy)*(√2 – 1) + (dx + dy – min(dx,dy))
      = (√2 – 1)*min(dx,dy) + (dx + dy) */
    float Astar::get_heuristic_cost(const NodePtr from, const NodePtr &to)
    {
        float dx = std::abs(to->get_coord_x() - from->get_coord_x());
        float dy = std::abs(to->get_coord_y() - from->get_coord_y());
        return (sqrt2_2 - 1) * std::min(dx, dy) + (dx + dy);
    }

    // 扩张邻节点
    void Astar::expansion_neighbors(const int &current_index, std::vector<int> &neighbor_index, std::vector<float> &neighbor_cost)
    {
        neighbor_index.clear();
        neighbor_cost.clear();

        int i = current_index % xs_; // 列号
        int j = current_index / xs_; // 行号

        if (j > 0) // 上
        {
            neighbor_index.emplace_back(current_index - xs_);
            neighbor_cost.emplace_back(1.0);
        }
        if (j < ys_ - 1) // 下
        {
            neighbor_index.emplace_back(current_index + xs_);
            neighbor_cost.emplace_back(1.0);
        }
        if (i > 0) //  左
        {
            neighbor_index.emplace_back(current_index - 1);
            neighbor_cost.emplace_back(1.0);
        }
        if (i < xs_ - 1) // 右
        {
            neighbor_index.emplace_back(current_index + 1);
            neighbor_cost.emplace_back(1.0);
        }
        if (motion_model_ == MotionModel::MOORE)
        {
            if (i > 0 && j > 0) // 左上
            {
                neighbor_index.emplace_back(current_index - xs_ - 1);
                neighbor_cost.emplace_back(sqrt2);
            }
            if (i > 0 && j < ys_ - 1) // 左下
            {
                neighbor_index.emplace_back(current_index + xs_ - 1);
                neighbor_cost.emplace_back(sqrt2);
            }
            if (i < xs_ - 1 && j < ys_ - 1) // 右下
            {
                neighbor_index.emplace_back(current_index + xs_ + 1);
                neighbor_cost.emplace_back(sqrt2);
            }
            if (i < xs_ - 1 && j > 0) // 右上
            {
                neighbor_index.emplace_back(current_index - xs_ + 1);
                neighbor_cost.emplace_back(sqrt2);
            }
        }
    }
    // 路径回溯
    bool Astar::backtrace_path(NodePtr goal, std::vector<Eigen::Vector2i> &path)
    {
        if (!goal)
        {
            return false;
        }
        NodePtr current_node = goal;
        while (current_node)
        {
            path.emplace_back(current_node->get_coord());
            current_node = current_node->parent_;
        }
        return !path.empty();
    }
    nav_msgs::msg::OccupancyGrid Astar::visualize()
    {
        nav_msgs::msg::OccupancyGrid map;
        map.header.frame_id = "map";
        map.info.height = costmap_->getSizeInCellsY();
        map.info.width = costmap_->getSizeInCellsX();
        map.info.resolution = costmap_->getResolution();
        map.info.origin.position.x = costmap_->getOriginX();
        map.info.origin.position.y = costmap_->getOriginY();
        map.info.origin.position.z = 0.0;
        map.info.origin.orientation.w = 1.0;
        map.data.assign(map.info.width * map.info.height, -1);

        for (size_t j = 0; j < map.info.height; ++j)
        {
            for (size_t i = 0; i < map.info.width; ++i)
            {
                int id = get_index(i, j);
                if (graph_[id].was_queued())
                {
                    map.data[id] = 60;
                }
                else if (graph_[id].was_visited())
                {
                    map.data[id] = 0;
                }
                else if (graph_[id].get_cost() == OCCUPIED)
                {
                    map.data[id] = 100;
                }
            }
        }
        return map;
    }
    // 核心 路径生成
    bool Astar::create_path(const Eigen::Vector2i &start, const Eigen::Vector2i &end, std::vector<Eigen::Vector2i> &path)
    {
        set_costmap(costmap_);

        int start_i = get_index(start);
        int end_i = get_index(end);

        // 边界检查
        if (start_i < 0 || start_i >= ns_ || end_i < 0 || end_i >= ns_)
        {
            std::cerr << "Astar: 起点或终点坐标不合法！start = " << start << " , end = " << end << std::endl;
            return false;
        }
        // 清空openlist
        clear_queue();

        NodePtr start_ptr = &(graph_[start_i]);
        NodePtr end_ptr = &(graph_[end_i]);
        // 起点加入openlist
        start_ptr->update(0.0, get_heuristic_cost(start_ptr, end_ptr)); // 计算start的g、h
        start_ptr->queued();
        open_list_.push(start_ptr);

        unsigned int iterations = 0; // 迭代数
        bool find_path = false;

        NodePtr current;
        while (!open_list_.empty())
        {
            iterations++;

            // 列返回f值最小的节点
            current = open_list_.top();
            open_list_.pop();
            // current->visited();
            if (current->was_visited())
            {
                continue;
            }

            if (current->get_index() == end_i)
            {
                std::cout << "找到终点！" << std::endl;
                find_path = true;
                break;
            }

            NodePtr tmp;
            std::vector<int> neighbor_index;
            std::vector<float> neighbor_cost;
            expansion_neighbors(current->get_index(), neighbor_index, neighbor_cost);

            for (size_t i = 0; i < neighbor_index.size(); ++i)
            {
                int t_i = neighbor_index[i];
                if (t_i < ns_ && t_i >= 0)
                {
                    tmp = &(graph_[t_i]);
                    if (tmp->get_cost() >= OCCUPIED || tmp->was_visited())
                    {
                        continue;
                    }
                    float tmp_g = current->G_x() + neighbor_cost[i];
                    if (tmp->was_unknown())
                    {
                        tmp->update(tmp_g, get_heuristic_cost(tmp, end_ptr));
                        tmp->parent_ = current;
                        open_list_.push(tmp);
                        tmp->queued();
                    }
                    else if (tmp_g < tmp->G_x())
                    {
                        // 发现更优路径，更新节点并重新加入队列
                        tmp->update(tmp_g, get_heuristic_cost(tmp, end_ptr));
                        tmp->parent_ = current;
                        open_list_.push(tmp); // 重新插入，旧节点将被跳过
                    }
                }
            }
        } // while
        if (find_path)
        {
            std::cout << "找到终点！迭代数：" << iterations << std::endl;
            std::vector<Eigen::Vector2i> path_reverse;

            if (!backtrace_path(end_ptr, path_reverse))
            {
                std::cout << "失败警告：找到路径，但是路径回溯失败！" << std::endl;
                return false;
            }
            else
            {
                path.clear();
                for (int i = path_reverse.size() - 1; i >= 0; i--)
                {
                    path.emplace_back(path_reverse[i]);
                }
            }
            visualize();
            return true;
        }
        else
        {
            std::cout << "没有找到路径！" << std::endl;
        }
        return false;
    }

} // namespace astar_planner